AWS Batch で並列演算ライブラリ Dask を利用したマルチノード並列ジョブを実行するワークショップに参加しました #AWSreInvent

AWS re:Invent 2023

#HPC

#AWS Batch

#AWS

大村保貴

2023.11.30

re:Invent 2023 の現地で開催された「Running high performance computing applications using AWS Batch」ワークショップに参加してきました。私が勉強になった点を中心に紹介します。

ワークショップの概要

タイトル

Running high performance computing applications using AWS Batch

概要

AWS Batch can run various types of workloads, such as machine learning and training LLMs, processing logs for compliance, and traditional high performance computing (HPC) applications like computational fluid dynamics. Join this workshop to learn how to run HPC and other workloads using AWS Batch. Learn AWS Batch fundamentals and then get hands-on experience with advanced AWS Batch features. Also learn how to package applications for different architectures (including AWS Graviton-based architectures), create fair-share scheduling policies to balance resources across different workloads, and visualize job placements to verify your scheduler configurations. You must bring your laptop to participate.

機械翻訳

AWS Batchは、機械学習やトレーニングLLM、コンプライアンスのためのログ処理、計算流体力学のような従来のハイパフォーマンスコンピューティング（HPC）アプリケーションなど、様々なタイプのワークロードを実行できます。このワークショップに参加して、AWS Batchを使用してHPCやその他のワークロードを実行する方法を学びましょう。AWS Batchの基礎を学び、AWS Batchの高度な機能を実際に体験します。また、異なるアーキテクチャ（AWS Gravitonベースのアーキテクチャを含む）向けにアプリケーションをパッケージ化する方法、異なるワークロード間でリソースのバランスを取るためのフェアシェアスケジューリングポリシーを作成する方法、スケジューラ構成を検証するためのジョブ配置を可視化する方法についても学びます。参加にはノートPCが必要です。

スピーカー

Maxime Hugues
Tom Burggraf

レベル

300

備考

ワークショップ内容は一般に公開されているリンクはありませんでした。（2023 年 11 月 28 日時点）

AWS Workshops

ワークショップの内容

ジョブ実行テスト用のコンテナイメージを作成して 3 種類のジョブ実行パターンを試すハンズオンでした。

コンテナイメージをビルド
ECR へ保存
AWS Batch の設定
ジョブ実行 3 パターン
- シングルジョブ
- 配列ジョブ
- マルチノード並列ジョブ

個人的に勉強になった配列ジョブと、マルチノード並列ジョブの実行についてまとめです。

配列ジョブ

stress-ngを利用して負荷をかけるイメージを作成しました。なにかの検証時のテストジョブに利用できそうなのでメモしておきます。

FROM public.ecr.aws/amazonlinux/amazonlinux:2
RUN yum -y update
RUN amazon-linux-extras install epel -y
RUN yum -y install stress-ng

### Build mktests.sh
RUN echo $'#!/bin/bash\n\
FILE=/stress-tests.txt\n\
rm $FILE 2>/dev/null\n\
COUNT=0\n\
for II in `stress-ng --cpu-method which 2>&1`\n\
do\n\
    if [ $COUNT -gt 5 ]; then\n\
        echo "--cpu 1 -t 120s --times --cpu-method $II" >> $FILE\n\
    fi\n\
    COUNT=`expr $COUNT + 1`\n\
done' >> /mktests.sh
RUN chmod 0744 /mktests.sh
RUN cat /mktests.sh
RUN echo $'#!/bin/bash\n\
/mktests.sh\n\
STRESS_ARGS=`sed -n $((AWS_BATCH_JOB_ARRAY_INDEX + 1))p /stress-tests.txt`\n
echo "Passing the following arguments to stress-ng: $STRESS_ARGS"\n\
/usr/bin/stress-ng $STRESS_ARGS' \n\ >> /docker-entrypoint.sh
RUN chmod 0744 /docker-entrypoint.sh
RUN cat /docker-entrypoint.sh
ENTRYPOINT ["/docker-entrypoint.sh"]

上記のコンテナイメージを利用して配列ジョブを実行しました。配列ジョブについてはワークショップの内容より、図が書いてある方がイメージしやすいと思いますので参考までに以下のブログを紹介します。

マルチノード並列ジョブ

私は OpenMPI を利用した並列ジョブを実行することはあります。本ワークショップでは Dask という Python 並列計算を行うためのライブラリを利用して並列ジョブを実行する点が目新しかったです。

画像引用: ワークショップのテキストより

マルチノード並列ジョブの AWS Batch のメインノード（EC2）で Dask scheduler が起動し、AWS Batch のワーカーノード（EC2）では Dask Worker が起動し並列分散で計算可能の様です。並列計算用のライブラリは OpneMPI か、 NVIDIA の CUDA で GPU 処理の経験しかなかったため興味深かかったです。とくに Python で書けるところが扱いやすそうなので今後検証で使えそうです。

Multi-node parallel jobs - AWS Batch

AWS Batch でマルチノード並列ジョブの実行は以下のブログと概ね同じです。今回は Dask で並列ジョブの管理しているだけの違いです。ただ、Dask のイメージ、テストジョブはワークショップ用に必要ファイルを一式まとめてダウンロードして実行するかたちでした。Cloud 9 上で実行しており、ローカルに退避し忘れたため欲しかったデータを失いました。

ジョブを実行するとメインノードが 1 台、ワーカーノードが 2 台起動しました。

ほどなくしてジョブは成功しました。

おわりに

目新しいものとしては Dask の存在を知りました。Dask の使い方を学ぶワークショップではないため、こちらのライブラリの使い方まではわかっていないです。マルチノード並列ジョブの検証する機会があれば利用してみようと思います。並列演算は C か C++ で書く印象があり、テストジョブくらいなら雰囲気で書いてなんとかしていたのですが、Python が使えるのであれば少し凝ったテストジョブも書けそうで検証がはかどりそうです。